#
#  Copyright 2009-2010 University of California
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

import os, sys
import exceptions
import copperhead.runtime.places as P
from copperhead.decorators import cutype, cushape, cubox
from copperhead.compiler.shapeinference import shapeof


_curdir = os.path.dirname(os.path.abspath(__file__))

#BLAS 2 operations
@cubox((_curdir, 'coblas.h'), prealloc=True)
@cutype("(Int, Int, Int, Float, [[Float]], [Float], Float, [Float]) -> [Float]")
@cushape(lambda transA, m, ku, alpha, A, X, beta, Y: (shapeof(Y), []))
def sgbmv(*args):
    pass

@cubox((_curdir, 'coblas.h'), prealloc=True)
@cutype("(Int, Float, [[Float]], [Float], Float, [Float]) -> [Float]")
@cushape(lambda transA, alpha, A, X, beta, Y: (shapeof(Y), []))
def sgemv(*args):
    pass

@cubox((_curdir, 'coblas.h'), prealloc=True)
@cutype("(Float, [Float], [Float], [[Float]]) -> [[Float]]")
@cushape(lambda alpha, X, Y, A: (shapeof(A), []))
def sger(*args):
    pass

@cubox((_curdir, 'coblas.h'), prealloc=True)
@cutype("(Int, Float, [[Float]], [Float], Float, [Float]) -> [Float]")
@cushape(lambda uplo, alpha, A, X, beta, Y: (shapeof(Y), []))
def ssymv(*args):
    pass

#BLAS 3 operations
@cubox((_curdir, 'coblas.h'), prealloc=True)
@cutype("(Int, Int, Float, [[Float]], [[Float]], Float, [[Float]]) -> [[Float]]")
@cushape(lambda transA, transB, alpha, A, B, beta, C: (shapeof(C), []))
def sgemm(*args):
    pass

try:
    from copperhead.runtime import nvcc_toolchain, host_toolchain
   
    current_path = os.path.dirname(os.path.abspath(__file__))
    
    nvcc_toolchain.add_library('cublas',
                               [], [], ['cublas'])
    host_toolchain.add_library('cublas',
                               [], [], ['cublas'])
    #blas2
    sgbmv.variants[P.here] = lambda *args:0
    sgemv.variants[P.here] = lambda *args:0
    sger.variants[P.here] = lambda *args:0
    ssymv.variants[P.here] = lambda *args:0
    sgbmv.variants[P.gpu0] = sgbmv
    sgemv.variants[P.gpu0] = sgemv
    sger.variants[P.gpu0] = sger
    ssymv.variants[P.gpu0] = ssymv

    #blas3
    sgemm.variants[P.here] = lambda *args:0
    sgemm.variants[P.gpu0] = sgemm

except ImportError as inst:
    # If the GPU place exists, we should report the error
    if hasattr(P, 'gpu0'):
        raise inst
